# engine
import torch
from torch.utils.data import DataLoader
from classes import BoltDataset, CONFIG, log
from models import PhysicsInformedNN
from gui import RealTimePlotter
import matplotlib.pyplot as plt
import os

class Trainer:
    def __init__(self):
        self.dataset = BoltDataset()
        self.loader = DataLoader(self.dataset, batch_size=CONFIG["batch_size"], shuffle=True)
        self.model = PhysicsInformedNN().to(CONFIG["device"])
        self.optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=CONFIG["lr"],
            weight_decay=CONFIG["weight_decay"]
        )
        self.best_loss = float("inf")
        self.plotter = RealTimePlotter() if plt.get_backend() == 'TkAgg' else None

    def train(self):
        log.info("开始训练，数据量：%d", len(self.dataset))
        try:
            for epoch in range(1, CONFIG["epochs"] + 1):
                total_loss = 0.0
                data_loss_sum = 0.0
                physics_loss_sum = 0.0
                
                for batch in self.loader:
                    V, F, A, E, g31, t, V_raw = [b.to(CONFIG["device"]) for b in batch]
                    
                    F_pred = self.model(V)
                    data_loss = torch.nn.functional.huber_loss(F_pred, F, delta=0.1)
                    
                    F_pred_raw = self.dataset.denormalize_F(F_pred.detach().cpu().numpy())
                    F_pred_raw = torch.tensor(F_pred_raw, device=CONFIG["device"])
                    
                    V_pred_physics = PhysicsInformedNN.physics_equation(F_pred_raw, A, E, g31, t)
                    physics_loss = torch.mean((V_raw - V_pred_physics) ** 2)
                    
                    loss = data_loss + 5 * physics_loss
                    
                    self.optimizer.zero_grad()
                    loss.backward()
                    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
                    self.optimizer.step()
                    
                    total_loss += loss.item() * len(V)
                    data_loss_sum += data_loss.item() * len(V)
                    physics_loss_sum += physics_loss.item() * len(V)
                
                avg_loss = total_loss / len(self.dataset)
                avg_data_loss = data_loss_sum / len(self.dataset)
                avg_physics_loss = physics_loss_sum / len(self.dataset)
                
                log.info(
                    "Epoch %4d | 总损失: %8.4e | 数据损失: %8.4e | 物理损失: %8.4e",
                    epoch, avg_loss, avg_data_loss, avg_physics_loss
                )
                
                # 实时绘图更新
                if self.plotter and epoch % CONFIG["plot_interval"] == 0:
                    with torch.no_grad():
                        batch = next(iter(self.loader))
                        V, F, A, E, g31, t, V_raw = [b.to(CONFIG["device"]) for b in batch]
                        F_pred = self.model(V)
                        F_pred_raw = self.dataset.denormalize_F(F_pred.detach().cpu().numpy())
                        V_pred_physics = PhysicsInformedNN.physics_equation(
                            torch.tensor(F_pred_raw, device=CONFIG["device"]), A, E, g31, t
                        )
                        self.plotter.update(
                            epoch, 
                            V_raw.cpu().numpy()[0], 
                            V_pred_physics.cpu().numpy()[0], 
                            F_pred_raw[0]
                        )
                
                # 保存模型
                if avg_loss < self.best_loss:
                    self.best_loss = avg_loss
                    self._save_checkpoint(epoch)
                    
        except Exception as e:
            log.error("训练过程中发生错误: %s", str(e))
            raise
        finally:
            if self.plotter:
                self.plotter.close() 

    def _save_checkpoint(self, epoch: int):
        state = {
            "epoch": epoch,
            "model_state_dict": self.model.state_dict(),
            "V_mean": self.dataset.V_mean,
            "V_std": self.dataset.V_std,
            "F_mean": self.dataset.F_mean,
            "F_std": self.dataset.F_std
        }
        torch.save(state, os.path.join(CONFIG["save_dir"], f"epoch_{epoch}.pth"))
        log.info("保存模型检查点,epoch: %d", epoch)